#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

// 矩阵维度配置（可修改）
#define M 1024    // 矩阵A的行数
#define N 512     // 矩阵A的列数/矩阵B的行数
#define K 2048    // 矩阵B的列数

// 块大小配置（可调整优化）
#define BLOCK_SIZE 32

// CUDA错误检查宏
#define CHECK(call) {                                                         \
    const cudaError_t error = call;                                            \
    if (error != cudaSuccess) {                                                \
        printf("Error: %s:%d, ", __FILE__, __LINE__);                         \
        printf("code:%d, reason: %s\n", error, cudaGetErrorString(error));    \
        exit(1);                                                               \
    }                                                                          \
}

// GPU核函数：矩阵乘法（基础版本）
__global__ void matrixMulKernel(float* A, float* B, float* C, int m, int n, int k) {
    // 计算全局行列索引
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < m && col < k) {
        float sum = 0.0f;
        for (int i = 0; i < n; ++i) {
            sum += A[row * n + i] * B[i * k + col];
        }
        C[row * k + col] = sum;
    }
}

// CPU参考实现：矩阵乘法（用于验证）
void matrixMulCPU(float* A, float* B, float* C, int m, int n, int k) {
    for (int row = 0; row < m; ++row) {
        for (int col = 0; col < k; ++col) {
            float sum = 0.0f;
            for (int i = 0; i < n; ++i) {
                sum += A[row * n + i] * B[i * k + col];
            }
            C[row * k + col] = sum;
        }
    }
}

// 初始化矩阵（随机值）
void initMatrix(float* mat, int size) {
    for (int i = 0; i < size; ++i) {
        mat[i] = (float)rand() / RAND_MAX;  // [0, 1]范围随机数
    }
}

// 验证结果（比较GPU和CPU计算结果）
bool verifyResult(float* C_ref, float* C_gpu, int size) {
    const float epsilon = 1e-4f;
    for (int i = 0; i < size; ++i) {
        if (fabs(C_ref[i] - C_gpu[i]) > epsilon) {
            printf("Mismatch at index %d: CPU=%.5f, GPU=%.5f\n", 
                  i, C_ref[i], C_gpu[i]);
            return false;
        }
    }
    return true;
}

int main() {
    // 分配主机内存
    size_t size_A = M * N * sizeof(float);
    size_t size_B = N * K * sizeof(float);
    size_t size_C = M * K * sizeof(float);
    
    float *h_A = (float*)malloc(size_A);
    float *h_B = (float*)malloc(size_B);
    float *h_C = (float*)malloc(size_C);
    float *h_C_ref = (float*)malloc(size_C);

    // 初始化输入矩阵
    srand(2023);
    initMatrix(h_A, M*N);
    initMatrix(h_B, N*K);

    // 分配设备内存
    float *d_A, *d_B, *d_C;
    CHECK(cudaMalloc(&d_A, size_A));
    CHECK(cudaMalloc(&d_B, size_B));
    CHECK(cudaMalloc(&d_C, size_C));

    // 拷贝数据到设备
    CHECK(cudaMemcpy(d_A, h_A, size_A, cudaMemcpyHostToDevice));
    CHECK(cudaMemcpy(d_B, h_B, size_B, cudaMemcpyHostToDevice));

    // 设置执行配置
    dim3 block(BLOCK_SIZE, BLOCK_SIZE);
    dim3 grid((K + block.x - 1) / block.x, (M + block.y - 1) / block.y);

    // 执行核函数并计时
    cudaEvent_t start, stop;
    CHECK(cudaEventCreate(&start));
    CHECK(cudaEventCreate(&stop));
    
    CHECK(cudaEventRecord(start));
    matrixMulKernel<<<grid, block>>>(d_A, d_B, d_C, M, N, K);
    CHECK(cudaEventRecord(stop));
    CHECK(cudaEventSynchronize(stop));
    
    float milliseconds = 0;
    CHECK(cudaEventElapsedTime(&milliseconds, start, stop));
    printf("GPU Time: %.3f ms\n", milliseconds);

    // 拷贝结果回主机
    CHECK(cudaMemcpy(h_C, d_C, size_C, cudaMemcpyDeviceToHost));

    // CPU计算参考结果
    matrixMulCPU(h_A, h_B, h_C_ref, M, N, K);

    // 验证结果
    if (verifyResult(h_C_ref, h_C, M*K)) {
        printf("Result Verification: PASSED\n");
    }

    // 释放资源
    free(h_A);
    free(h_B);
    free(h_C);
    free(h_C_ref);
    CHECK(cudaFree(d_A));
    CHECK(cudaFree(d_B));
    CHECK(cudaFree(d_C));
    CHECK(cudaEventDestroy(start));
    CHECK(cudaEventDestroy(stop));

    return 0;
}
